/*
 * This is written by Lulu Han.
 * E-mail: locomotive_crypto@163.com
 */

#include <iostream>
#include <exception>
#include "sec_cmp.h"
#include "timer.h"
#include "net_client.h"
#include "net_server.h"

void test_cmp_for_correctness(long bitlens);

void test_cmp_for_performance(long bits, long cnt);

void test_cmp_by_network_SA(long bits, long cnt, size_t port, std::string ipaddr);

void test_cmp_by_network_SB(size_t port);

int main(int argc, char* argv[]){
	
	std::cout << "----------------------------------------------------------" << std::endl;
	std::cout << "-----------------------<Usage>----------------------------" << std::endl;
	std::cout << "-----1) Run for correctness: ./program--------------------" << std::endl;
	std::cout << "-----2) Run for performance: ./program cnt----------------" << std::endl;
	std::cout << "-----3) Run as SB: ./program port ------------------------" << std::endl;
	std::cout << "-----4) Run as SA: ./program port ipaddr------------------" << std::endl;
	std::cout << "----------------------------------------------------------" << std::endl << std::endl;
	
	if(argc == 1){
		test_cmp_for_correctness(200);
	}
	else if(argc == 2){
		for(size_t i = 0; i < 5; ++i)
			test_cmp_for_performance(512, std::atoi(argv[1]));
		std::cout << std::endl;
		
		for(size_t i = 0; i < 5; ++i)
			test_cmp_for_performance(1024, std::atoi(argv[1]));
		std::cout << std::endl;
		
		for(size_t i = 0; i < 5; ++i)
			test_cmp_for_performance(1536, std::atoi(argv[1]));
		std::cout << std::endl;
	}
	else if(argc == 3){
		if(std::string(argv[1]) == std::string("sb") || std::string(argv[1]) == std::string("SB")){
			test_cmp_by_network_SB(std::atoi(argv[2]));
		}
	}
	else if(argc == 5){
		if(std::string(argv[1]) == std::string("sa") || std::string(argv[1]) == std::string("SA")){
			//test_cmp_by_network_SA(512, std::atoi(argv[4]), std::atoi(argv[2]), std::string(argv[3]));
			
			//test_cmp_by_network_SA(1024, std::atoi(argv[4]), std::atoi(argv[2]), std::string(argv[3]));
			
			test_cmp_by_network_SA(1536, std::atoi(argv[4]), std::atoi(argv[2]), std::string(argv[3]));
		}	
	}
	return 0;
}

void test_cmp_for_correctness(long bitlens){
	PrivateKey prikey;
	PublicKey pubkey;
	
	Paillier::GenKeyPair(prikey, pubkey, bitlens);
	
	Decryptor dec(prikey);
	Encryptor enc(pubkey);
	
	NTL::ZZ x, y;
	
	x = NTL::RandomLen_ZZ(16);
	y = NTL::RandomLen_ZZ(16);
	
	std::cout << "x = " << x << std::endl;
	std::cout << "y = " << y << std::endl;
	
	if(x >= y){
		std::cout << "x >= y: " << (x >= y) << std::endl;
	}
	else{
		std::cout << "x < y: " << (x >= y) << std::endl;
	}
	
	SecCmp::CloudA CA(prikey);
	SecCmp::CloudB CB(pubkey);
	
	CA.set_nvlen(bitlens*2-1, 32);
	CB.set_nvlen(bitlens*2-1, 32);
	
	NTL::ZZ cx, cy;
	enc.Encrypt(x, cx);
	enc.Encrypt(y, cy);
	
	NTL::ZZ csub;
	CB.calc_sub(cx, cy, csub);
	
	NTL::ZZ cbeta;
	CA.dec_sub(csub, cbeta);
	
	NTL::ZZ nca, ncb;
	CB.rerandom_for_mul(cbeta, nca, ncb);
	
	NTL::ZZ cmul;
	CA.dec_reenc(nca, ncb, cmul);
	
	NTL::ZZ ccmp;
	CB.gen_res(cmul, ccmp);
	
	NTL::ZZ cmp;
	dec.Decrypt(ccmp, cmp);
	
	std::cout << "cmp = " << cmp << std::endl;
}

void test_cmp_for_performance(long bits, long cnt){
	std::cout << "The key length for this test is: " << 2*bits << std::endl;
	 
	std::vector<Paillier::Encryptor> vencryptor;
	std::vector<Paillier::Decryptor> vdecryptor;
	
	std::vector<SecCmp::CloudA> vCA;
	std::vector<SecCmp::CloudB> vCB;
	
	for(long i = 0; i < cnt; ++i){
		Paillier::PrivateKey prikey;
		Paillier::PublicKey pubkey;
	
		// Generate random key pair
		Paillier::GenKeyPair(prikey, pubkey, bits);
	
		// Define Encryptor and Decryptor objects
		Paillier::Encryptor encryptor(pubkey);
		Paillier::Decryptor decryptor(prikey);
		
		vencryptor.push_back(encryptor);
		vdecryptor.push_back(decryptor);
		
		SecCmp::CloudA CA(prikey);
		SecCmp::CloudB CB(pubkey);
		
		vCA.push_back(CA);
		vCB.push_back(CB);
	}
	
	std::vector<NTL::ZZ> va, vb;
	
	NTL::ZZ r;
	for(long i = 0; i < cnt; ++i){
		r = NTL::RandomLen_ZZ(30);
		va.push_back(r);
		
		r = NTL::RandomLen_ZZ(30);
		vb.push_back(r);
	}
	
	std::vector<NTL::ZZ> vca, vcb;
	NTL::ZZ c;
	for(long i = 0; i < cnt; ++i){
		vencryptor[i].Encrypt(va[i], c);
		vca.push_back(c);
		
		vencryptor[i].Encrypt(vb[i], c);
		vcb.push_back(c);
	}
	
	for(long i = 0; i < cnt; ++i){
		vCA[i].set_nvlen(bits*2-1, 32);
		vCB[i].set_nvlen(bits*2-1, 32);
	}
	
	
	time_point<high_resolution_clock> start;	
	start = now_time();
	for(long i = 0; i < cnt; ++i){
		NTL::ZZ csub;
		vCB[i].calc_sub(vca[i], vcb[i], csub);
	
		NTL::ZZ cbeta;
		vCA[i].dec_sub(csub, cbeta);
	
		NTL::ZZ nca, ncb;
		vCB[i].rerandom_for_mul(cbeta, nca, ncb);
	
		NTL::ZZ cmul;
		vCA[i].dec_reenc(nca, ncb, cmul);
	
		NTL::ZZ ccmp;
		vCB[i].gen_res(cmul, ccmp);
	
		/*
		NTL::ZZ pcmp;
		vdecryptor[i].Decrypt(ccmp, pcmp);
	
		std::cout << "(va[i] >= vb[i]): " << (va[i] >= vb[i]) << std::endl;
		std::cout << "pcmp: " << pcmp << std::endl;
		
		if((va[i] >= vb[i]) != (pcmp == 1)){
			throw std::runtime_error("Failed to perform SecCmp...\n");
		}*/
	}
	double tm = time_from(start);
	std::cout << "The time for "<< cnt << " SecCmp is : " << tm << std::endl;

}

void test_cmp_by_network_SA(long bits, long cnt, size_t port, std::string ipaddr){
	std::cout << "The key length for this test is: " << 2*bits << std::endl;
	
	NetClient client(ipaddr, port);
	client.connect();
	
	client.send(cnt);
	client.send(bits);
	
	std::vector<Paillier::Decryptor> vdecryptor;
	
	std::vector<SecCmp::CloudA> vCA;
	
	for(long i = 0; i < cnt; ++i){
		Paillier::PrivateKey prikey;
		Paillier::PublicKey pubkey;
	
		// Generate random key pair
		Paillier::GenKeyPair(prikey, pubkey, bits);
		
		client.send(pubkey.GetN());
		
		// Define Encryptor and Decryptor objects
		Paillier::Decryptor decryptor(prikey);
		
		vdecryptor.push_back(decryptor);
		
		SecCmp::CloudA CA(prikey);
		
		vCA.push_back(CA);
	}
	
	for(long i = 0; i < cnt; ++i){
		vCA[i].set_nvlen(bits*2-1, 32);
	}
	
	std::cout << "Please count the network traffics from here..." << std::endl;
	getchar();
	for(long i = 0; i < cnt; ++i){
		NTL::ZZ csub;
		client.receive(csub);
	
		NTL::ZZ cbeta;
		vCA[i].dec_sub(csub, cbeta);
		client.send(cbeta);
		
		NTL::ZZ nca, ncb;
		client.receive(nca);
		client.receive(ncb);
	
		NTL::ZZ cmul;
		vCA[i].dec_reenc(nca, ncb, cmul);
		client.send(cmul);
		
		//NTL::ZZ ccmp;
		//client.receive(ccmp);
	
		
		//NTL::ZZ pcmp;
		//vdecryptor[i].Decrypt(ccmp, pcmp);
		
		//std::cout << "pcmp: " << pcmp << std::endl;
	}
	
	client.close();
}

void test_cmp_by_network_SB(size_t port){
	NetServer server(port); 
	server.accept();
	
	long cnt;
	server.receive(cnt);
	
	long bits;
	server.receive(bits);
	
	std::cout << "The key length for this test is: " << 2*bits << std::endl;
	
	std::vector<Paillier::Encryptor> vencryptor;
	std::vector<SecCmp::CloudB> vCB;
	
	NTL::ZZ N;
	for(long i = 0; i < cnt; ++i){
		server.receive(N);
		
		Paillier::PublicKey pubkey(N);
		SecCmp::CloudB CB(pubkey);
		
		Paillier::Encryptor enc(pubkey);
		vencryptor.push_back(enc);
		
		vCB.push_back(CB);
	}
	
	std::vector<NTL::ZZ> va, vb;
	
	NTL::ZZ r;
	for(long i = 0; i < cnt; ++i){
		r = NTL::RandomLen_ZZ(30);
		va.push_back(r);
		
		r = NTL::RandomLen_ZZ(30);
		vb.push_back(r);
	}
	
	std::vector<NTL::ZZ> vca, vcb;
	NTL::ZZ c;
	for(long i = 0; i < cnt; ++i){
		vencryptor[i].Encrypt(va[i], c);
		vca.push_back(c);
		
		vencryptor[i].Encrypt(vb[i], c);
		vcb.push_back(c);
	}
	
	for(long i = 0; i < cnt; ++i){
		vCB[i].set_nvlen(bits*2-1, 32);
	}
	
	std::cout << "Please count the network traffics from here..." << std::endl;
	getchar();
	for(long i = 0; i < cnt; ++i){
		NTL::ZZ csub;
		vCB[i].calc_sub(vca[i], vcb[i], csub);
		server.send(csub);
	
		NTL::ZZ cbeta;
		server.receive(cbeta);
		
		NTL::ZZ nca, ncb;
		vCB[i].rerandom_for_mul(cbeta, nca, ncb);
		server.send(nca);
		server.send(ncb);
		
		NTL::ZZ cmul;
		server.receive(cmul);
	
		NTL::ZZ ccmp;
		vCB[i].gen_res(cmul, ccmp);
		//server.send(ccmp);
		
		//std::cout << "(va[i] >= vb[i]): " << (va[i] >= vb[i]) << std::endl;
	}
	
	server.close();
}
